#ifndef _MULTI_CAVITY_MATRIX_BUILDER_H_
#define _MULTI_CAVITY_MATRIX_BUILDER_H_

#include <iostream>
#include <vector>
#include <string>
#include <memory>
#include <map>
#include "pcg_def.h"
#include "tensor_mat_multi_cav.h"

// Debugging macros
#define DEBUG_LEVEL 1  // 0=none, 1=errors, 2=warnings, 3=info, 4=detailed
#define DEBUG_LOG(level, msg) if (level <= DEBUG_LEVEL) { std::cout << "[DEBUG][L" << level << "] " << msg << std::endl; }
#define DEBUG_ERROR(msg) DEBUG_LOG(1, "ERROR: " << msg)
#define DEBUG_WARN(msg) DEBUG_LOG(2, "WARNING: " << msg)
#define DEBUG_INFO(msg) DEBUG_LOG(3, "INFO: " << msg)
#define DEBUG_DETAIL(msg) DEBUG_LOG(4, "DETAIL: " << msg)

namespace MultiCavityMatrix {

// Structure to represent a 3D grid
struct Grid3D {
    int imt;  // Points in x direction
    int jmt;  // Points in y direction
    int kmt;  // Points in z direction
    
    // Physical domain extents
    double x_min, x_max;
    double y_min, y_max;
    double z_min, z_max;
    
    Grid3D(int i, int j, int k, 
          double x0 = 0.0, double x1 = 1.0, 
          double y0 = 0.0, double y1 = 1.0, 
          double z0 = 0.0, double z1 = 1.0)
        : imt(i), jmt(j), kmt(k),
          x_min(x0), x_max(x1),
          y_min(y0), y_max(y1),
          z_min(z0), z_max(z1) {}
    
    int total_points() const { return imt * jmt * kmt; }
    
    // Alias for total_points() to match the usage in to_csr method
    int size() const {
        return total_points();
    }
    
    // Debug information
    void print_info() const;
};

// Structure to represent interface between cavities
struct CavityInterface {
    int src_cavity_id;  // Source cavity
    int dst_cavity_id;  // Destination cavity
    
    // Interface region in source cavity (inclusive)
    int src_i_start, src_i_end;
    int src_j_start, src_j_end;
    int src_k_start, src_k_end;
    
    // Starting point in destination cavity
    int dst_i_start;
    int dst_j_start;
    int dst_k_start;
    
    // Constructor
    CavityInterface(
        int src_id, int dst_id,
        int si_start, int si_end, int sj_start, int sj_end, int sk_start, int sk_end,
        int di_start, int dj_start, int dk_start)
        : src_cavity_id(src_id), dst_cavity_id(dst_id),
          src_i_start(si_start), src_i_end(si_end),
          src_j_start(sj_start), src_j_end(sj_end),
          src_k_start(sk_start), src_k_end(sk_end),
          dst_i_start(di_start), dst_j_start(dj_start), dst_k_start(dk_start) {}
    
    // Debug information
    void print_info() const;
};

// Method to use for discretization
enum class DiscretizationMethod {
    FDM,         // Finite Difference Method (19-point stencil)
    SEM,         // Spectral Element Method
    SPECTRAL     // Spectral Method
};

// Class to build coefficient matrices for Poisson equation in multiple cavities
class MultiCavityMatrixBuilder {
private:
    // Cavity grids
    std::vector<Grid3D> cavity_grids;
    
    // Interfaces between cavities
    std::vector<CavityInterface> interfaces;
    
    // Total number of points across all cavities
    int total_grid_points;
    
    // Method used for discretization
    DiscretizationMethod method;
    
    // Number of diagonals in the matrix
    int num_diagonals;
    
    // Map for point indexing (global to cavity+local)
    std::map<int, std::pair<int, int>> global_to_local_map;
    
    // Store cavity offsets for global indexing
    std::vector<int> cavity_offsets;
    
    // Initialize indexing for points
    void initialize_indexing();
    
    // Build differential operator for a single cavity
    void build_cavity_operator(int cavity_id, float* tensor_data);
    
    // Connect cavities at interfaces
    void connect_cavities(float* tensor_data);
    
    // Helper to apply boundary conditions
    void apply_boundary_conditions(float* tensor_data);
    
    // Debug helpers
    void validate_interfaces() const;
    void print_tensor_stats(const float* tensor_data) const;
    
    // Helper functions for the to_csr conversion
    // Get neighbor indices for a given diagonal
    void get_neighbor_indices(int diagonal, int& i, int& j, int& k, int& cavity_id, const Grid3D& grid) const {
        // Map diagonal index to relative coordinates
        int di = 0, dj = 0, dk = 0;
        
        // Unused parameters to prevent warnings
        (void)cavity_id;
        (void)grid;
        
        // For FDM 19-point stencil
        switch (diagonal) {
            case 0:  di=0;  dj=0;  dk=0;  break; // Center point (i,j,k)
            case 1:  di=1;  dj=0;  dk=0;  break; // i+1,j,k
            case 2:  di=-1; dj=0;  dk=0;  break; // i-1,j,k
            case 3:  di=0;  dj=1;  dk=0;  break; // i,j+1,k
            case 4:  di=0;  dj=-1; dk=0;  break; // i,j-1,k
            case 5:  di=0;  dj=0;  dk=1;  break; // i,j,k+1
            case 6:  di=0;  dj=0;  dk=-1; break; // i,j,k-1
            case 7:  di=1;  dj=1;  dk=0;  break; // i+1,j+1,k
            case 8:  di=1;  dj=-1; dk=0;  break; // i+1,j-1,k
            case 9:  di=-1; dj=1;  dk=0;  break; // i-1,j+1,k
            case 10: di=-1; dj=-1; dk=0;  break; // i-1,j-1,k
            case 11: di=1;  dj=0;  dk=1;  break; // i+1,j,k+1
            case 12: di=1;  dj=0;  dk=-1; break; // i+1,j,k-1
            case 13: di=-1; dj=0;  dk=1;  break; // i-1,j,k+1
            case 14: di=-1; dj=0;  dk=-1; break; // i-1,j,k-1
            case 15: di=0;  dj=1;  dk=1;  break; // i,j+1,k+1
            case 16: di=0;  dj=1;  dk=-1; break; // i,j+1,k-1
            case 17: di=0;  dj=-1; dk=1;  break; // i,j-1,k+1
            case 18: di=0;  dj=-1; dk=-1; break; // i,j-1,k-1
            default: break;
        }
        
        // Apply relative coordinates to get neighbor
        i += di;
        j += dj;
        k += dk;
    }
    
    // Check if indices are outside grid boundaries
    bool is_boundary(int i, int j, int k, const Grid3D& grid) const {
        return (i < 0 || i >= grid.imt || 
                j < 0 || j >= grid.jmt || 
                k < 0 || k >= grid.kmt);
    }
    
    // Check if a point is a valid interface point in the specified direction
    bool is_valid_interface_point(int cavity_id, int i, int j, int k, int diagonal) const {
        // Get neighbor direction
        int ni = i, nj = j, nk = k;
        int ncavity = cavity_id;
        get_neighbor_indices(diagonal, ni, nj, nk, ncavity, cavity_grids[cavity_id]);
        
        // If neighbor is within bounds, not an interface point
        if (!is_boundary(ni, nj, nk, cavity_grids[cavity_id])) {
            return false;
        }
        
        // Check each interface to see if this point is on it
        for (const auto& interface : interfaces) {
            if (interface.src_cavity_id == cavity_id) {
                // Check if the point is in the source region of this interface
                if (i >= interface.src_i_start && i <= interface.src_i_end &&
                    j >= interface.src_j_start && j <= interface.src_j_end &&
                    k >= interface.src_k_start && k <= interface.src_k_end) {
                    // Check if the neighbor direction points to the destination cavity
                    // This depends on the interface orientation
                    // For simplicity, just check if the interface is on the boundary
                    // where the neighbor would be
                    if ((i == interface.src_i_end && ni > i) ||
                        (i == interface.src_i_start && ni < i) ||
                        (j == interface.src_j_end && nj > j) ||
                        (j == interface.src_j_start && nj < j) ||
                        (k == interface.src_k_end && nk > k) ||
                        (k == interface.src_k_start && nk < k)) {
                        return true;
                    }
                }
            }
            
            if (interface.dst_cavity_id == cavity_id) {
                // Calculate relative position in destination region
                int rel_i = i - interface.dst_i_start;
                int rel_j = j - interface.dst_j_start;
                int rel_k = k - interface.dst_k_start;
                
                // Check if the point is in the destination region
                if (rel_i >= 0 && rel_i <= (interface.src_i_end - interface.src_i_start) &&
                    rel_j >= 0 && rel_j <= (interface.src_j_end - interface.src_j_start) &&
                    rel_k >= 0 && rel_k <= (interface.src_k_end - interface.src_k_start)) {
                    
                    // Check if the neighbor direction points to the source cavity
                    // This depends on the interface orientation
                    if ((i == interface.dst_i_start && ni < i) ||
                        (i == interface.dst_i_start + (interface.src_i_end - interface.src_i_start) && ni > i) ||
                        (j == interface.dst_j_start && nj < j) ||
                        (j == interface.dst_j_start + (interface.src_j_end - interface.src_j_start) && nj > j) ||
                        (k == interface.dst_k_start && nk < k) ||
                        (k == interface.dst_k_start + (interface.src_k_end - interface.src_k_start) && nk > k)) {
                        return true;
                    }
                }
            }
        }
        
        return false;
    }
    
    // Get global index of interface connection point
    int get_interface_connection(int cavity_id, int i, int j, int k, int diagonal) const {
        // Get neighbor direction
        int ni = i, nj = j, nk = k;
        int ncavity = cavity_id;
        get_neighbor_indices(diagonal, ni, nj, nk, ncavity, cavity_grids[cavity_id]);
        
        // Check each interface to find the connection
        for (const auto& interface : interfaces) {
            if (interface.src_cavity_id == cavity_id) {
                // Check if the point is in the source region of this interface
                if (i >= interface.src_i_start && i <= interface.src_i_end &&
                    j >= interface.src_j_start && j <= interface.src_j_end &&
                    k >= interface.src_k_start && k <= interface.src_k_end) {
                    
                    // Check if the neighbor direction points to the destination cavity
                    if ((i == interface.src_i_end && ni > i) ||
                        (i == interface.src_i_start && ni < i) ||
                        (j == interface.src_j_end && nj > j) ||
                        (j == interface.src_j_start && nj < j) ||
                        (k == interface.src_k_end && nk > k) ||
                        (k == interface.src_k_start && nk < k)) {
                        
                        // Map from source to destination coordinates
                        int dst_i = interface.dst_i_start + (i - interface.src_i_start);
                        int dst_j = interface.dst_j_start + (j - interface.src_j_start);
                        int dst_k = interface.dst_k_start + (k - interface.src_k_start);
                        
                        // Return global index in destination cavity
                        return local_to_global_3d(interface.dst_cavity_id, dst_i, dst_j, dst_k);
                    }
                }
            }
            
            if (interface.dst_cavity_id == cavity_id) {
                // Calculate relative position in destination region
                int rel_i = i - interface.dst_i_start;
                int rel_j = j - interface.dst_j_start;
                int rel_k = k - interface.dst_k_start;
                
                // Check if the point is in the destination region
                if (rel_i >= 0 && rel_i <= (interface.src_i_end - interface.src_i_start) &&
                    rel_j >= 0 && rel_j <= (interface.src_j_end - interface.src_j_start) &&
                    rel_k >= 0 && rel_k <= (interface.src_k_end - interface.src_k_start)) {
                    
                    // Check if the neighbor direction points to the source cavity
                    if ((i == interface.dst_i_start && ni < i) ||
                        (i == interface.dst_i_start + (interface.src_i_end - interface.src_i_start) && ni > i) ||
                        (j == interface.dst_j_start && nj < j) ||
                        (j == interface.dst_j_start + (interface.src_j_end - interface.src_j_start) && nj > j) ||
                        (k == interface.dst_k_start && nk < k) ||
                        (k == interface.dst_k_start + (interface.src_k_end - interface.src_k_start) && nk > k)) {
                        
                        // Map from destination to source coordinates
                        int src_i = interface.src_i_start + rel_i;
                        int src_j = interface.src_j_start + rel_j;
                        int src_k = interface.src_k_start + rel_k;
                        
                        // Return global index in source cavity
                        return local_to_global_3d(interface.src_cavity_id, src_i, src_j, src_k);
                    }
                }
            }
        }
        
        // No interface connection found
        return -1;
    }
    
    // Convert 3D indices to linear local index
    int local_index_3d(int i, int j, int k, const Grid3D& grid) const {
        return i + j * grid.imt + k * grid.imt * grid.jmt;
    }
    
public:
    MultiCavityMatrixBuilder(DiscretizationMethod method = DiscretizationMethod::FDM);
    ~MultiCavityMatrixBuilder();
    
    // Add a new cavity to the system
    int add_cavity(int imt, int jmt, int kmt, 
                 double x_min = 0.0, double x_max = 1.0,
                 double y_min = 0.0, double y_max = 1.0,
                 double z_min = 0.0, double z_max = 1.0);
    
    // Add an interface between cavities
    void add_interface(
        int src_cavity_id, int dst_cavity_id,
        int src_i_start, int src_i_end, int src_j_start, int src_j_end, int src_k_start, int src_k_end,
        int dst_i_start, int dst_j_start, int dst_k_start);
    
    // Build the multi-cavity coefficient matrix
    // Returns a 5D tensor as a flat array: [nct][m][imt][jmt][kmt]
    float* build_coefficient_matrix();
    
    // Convert global index to cavity + local indices
    void global_to_local(int global_idx, int& cavity_id, int& local_idx) const;
    void global_to_local_3d(int global_idx, int& cavity_id, int& i, int& j, int& k) const;
    
    // Convert cavity + local indices to global index
    int local_to_global(int cavity_id, int local_idx) const;
    int local_to_global_3d(int cavity_id, int i, int j, int k) const;
    
    // Matrix-vector multiplication with the coefficient matrix
    void matrix_vector_multiply(const float* tensor_data, const float* x, float* result) const;
    
    // Convert the tensor format to CSR format
    void to_csr(const float* tensor_data, CsrMatrix& csr_matrix) const;
    
    // Helper functions for CSR conversion
    int get_column_index(int global_idx, int stencil_idx) const;
    int get_stencil_column(int row, int stencil_idx, int cavity_id, 
                          int i, int j, int k, const Grid3D& grid) const;
    int get_spectral_column(int row, int stencil_idx, int cavity_id,
                           int i, int j, int k, const Grid3D& grid) const;
    float get_stencil_value(const float* tensor_data, int global_idx, int stencil_idx) const;
    
    // Get the number of diagonal elements
    int get_num_diagonals() const { return num_diagonals; }
    
    // Get total number of points
    int get_total_grid_points() const { return total_grid_points; }
    
    // Get number of cavities
    int get_num_cavities() const { return cavity_grids.size(); }
    
    // Print debug information
    void print_info() const;
    
    // Get dimensions of a specific cavity
    const Grid3D& get_cavity_dimensions(int cavity_id) const {
        if (cavity_id < 0 || cavity_id >= static_cast<int>(cavity_grids.size())) {
            DEBUG_ERROR("Invalid cavity ID: " << cavity_id);
            static Grid3D dummy(0, 0, 0);
            return dummy;
        }
        return cavity_grids[cavity_id];
    }
    
    // Check if a point is on an interface
    bool is_interface_point(int cavity_id, int i, int j, int k) const {
        if (cavity_id < 0 || cavity_id >= static_cast<int>(cavity_grids.size())) {
            return false;
        }
        
        // Check each interface
        for (const auto& interface : interfaces) {
            if (interface.src_cavity_id == cavity_id) {
                // Check if point is in the source region of this interface
                if (i >= interface.src_i_start && i <= interface.src_i_end &&
                    j >= interface.src_j_start && j <= interface.src_j_end &&
                    k >= interface.src_k_start && k <= interface.src_k_end) {
                    return true;
                }
            }
            
            if (interface.dst_cavity_id == cavity_id) {
                // Check if point is in the destination region of this interface
                int rel_i = i - interface.dst_i_start;
                int rel_j = j - interface.dst_j_start;
                int rel_k = k - interface.dst_k_start;
                
                if (rel_i >= 0 && rel_i <= (interface.src_i_end - interface.src_i_start) &&
                    rel_j >= 0 && rel_j <= (interface.src_j_end - interface.src_j_start) &&
                    rel_k >= 0 && rel_k <= (interface.src_k_end - interface.src_k_start)) {
                    return true;
                }
            }
        }
        
        return false;
    }
    
    // Get the cavity offset
    int get_cavity_offset(int cavity_id) const { return cavity_offsets[cavity_id]; }
};

// Utility functions
void dump_tensor_to_file(const float* tensor_data, int num_cavities, int num_diags, 
                         const std::vector<Grid3D>& grids, const std::string& filename);

// Return a human-readable method name
std::string method_to_string(DiscretizationMethod method);

} // namespace MultiCavityMatrix

#endif // _MULTI_CAVITY_MATRIX_BUILDER_H_ 